from pdf2image import convert_from_path
from PIL import Image, ImageDraw, ImageFont
from cnocr import CnOcr
from tqdm import tqdm
import shutil
import glob
import json
import os


def draw_text(src_img_path,dst_img_path):
    color = (255, 0, 0)
    outline="red"
    width=1
    ocr = CnOcr() 
    data = ocr.ocr(src_img_path)
    image = Image.open(src_img_path)
    draw = ImageDraw.Draw(image)
    font = ImageFont.truetype(r"C:\Windows\Fonts\STFANGSO.TTF", 17)
    for item in data:
        text = item['text']
        position = item['position']
        # bbox = draw.textbbox(
        #     tuple(position[0]), 
        #     text, 
        #     font=font, 
        #     spacing=0, 
        #     anchor='mm',
        #     align='center')
        x = int(position[0][0]+(position[2][0] - position[0][0])/2)
        y = int(position[0][1]+(position[2][1] - position[0][1])/2)
        draw.text(
            (x,y),
            text, 
            font=font, 
            anchor="mm", # align="center", 
            spacing=0,
            fill=color,
            stroke_width=1, 
            stroke_fill=color)
        xy = (
            item['position'][0][0], 
            item['position'][0][1], 
            item['position'][2][0], 
            item['position'][2][1])
        draw.polygon(points, outline='black', fill='blue')
        # draw.rectangle(xy=xy,fill=None,outline=outline,width=width)
    image.save(dst_img_path)
    return


def get_pdf(pdf_dir,save_dir):
    for name in tqdm(os.listdir(pdf_dir)):
        pdf_path=os.path.join(pdf_dir,name)
        dir_name=os.path.splitext(name)[0]
        main(pdf_path,save_dir,dir_name)
    return


def main(pdf_path,save_dir,dir_name):
    color = (255, 0, 0)
    outline="blue"
    width=2
    ttf_path=r"/data/fengyy/dateset/STKAITI.TTF"
    font = ImageFont.truetype(ttf_path, 28)
    # font = ImageFont.load_default()
    ocr = CnOcr()  # 所有参数都使用默认值
    dir_path = os.path.join(save_dir,dir_name)
    if not os.path.exists(dir_path):
        os.makedirs(dir_path)
    else:
        shutil.rmtree(dir_path)
        os.makedirs(dir_path)
    convert_from_path(pdf_path, fmt='jpg',output_folder=dir_path)
    paths = sorted(glob.glob(f"{dir_path}/*-*.jpg"))

    for i,img_path in enumerate(paths):
        
        new_path = os.path.join(dir_path,f"origin_{dir_name}_00{i+1}.jpg")
        os.rename(img_path,new_path)
        
        original_image = Image.open(new_path)
        
        width, height = original_image.size
        image = Image.new('RGB', (width, height), 'white')
        
        draw = ImageDraw.Draw(image)
        data = ocr.ocr(new_path)
        for j,item in enumerate(data):
            position = item['position']
            num_text = f"{j}"

            bbox = draw.textbbox(
                position[0], 
                num_text, 
                font=font, 
                spacing=0, 
                anchor='mm',
                align='center')
            point = (
                position[0][0]-(bbox[2]-bbox[0])/2,
                position[0][1]+(bbox[3]-bbox[1]))
            draw.text(
                point,
                num_text, 
                font=font, 
                anchor="mm", # align="center", 
                spacing=0,
                fill='green',
                stroke_width=1, 
                stroke_fill='green')
            
            
            x = int(position[0][0]+(position[2][0] - position[0][0])/2)
            y = int(position[0][1]+(position[2][1] - position[0][1])/2)
            draw.text(
                (x,y),
                item['text'], 
                font=font, 
                anchor="mm", # align="center", 
                spacing=0,
                fill=color,
                stroke_width=1, 
                stroke_fill=color)
   
            draw.polygon(
                position,
                outline="blue",
                fill=None,
                width=1)
            item['position']=item['position'].tolist()
            
        draw_path = os.path.join(dir_path,f"{dir_name}_00{i+1}.jpg")
        image.save(draw_path)
         
        json_name = f"{dir_name}_00{i+1}.json"
        json_path = os.path.join(dir_path,json_name).encode('utf-8')
        with open(json_path,"w",encoding='utf-8') as file_handle:
            json.dump(data,file_handle,ensure_ascii=False,indent=4)
        # os.remove(img_path)
            
    return


if __name__ == '__main__':
    # pdf_path=r"/data/fengyy/dateset/temp/1.农村土地承包经营权登记簿/模板一/1.农村土地承包经营权登记簿.pdf"
    # save_dir=r"/data/fengyy/dateset/temp/1.农村土地承包经营权登记簿/模板一"
    # dir_name = "模板一"
    # pdf_path=r"/data/fengyy/dateset/temp/1.农村土地承包经营权登记簿/模板二/4-0002.pdf"
    # pdf_path=r"/data/fengyy/dateset/temp/1.农村土地承包经营权登记簿/模板二/登记簿_吴推良_360622204200010001J_.pdf"
    # pdf_path=r"/data/fengyy/dateset/temp/1.农村土地承包经营权登记簿/模板二/登记簿.pdf"
    # save_dir=r"/data/fengyy/dateset/temp/1.农村土地承包经营权登记簿/模板二/模板二"
    # dir_name = "模板二_003"
    # main(pdf_path,save_dir,dir_name)
    
    # pdf_dir=r"/data/fengyy/dateset/temp/0429/农村土地承包经营权登记簿/模板二(1)/模板二_001"
    # save_dir=r"/data/fengyy/dateset/temp/0429/result/模板二_001"
    # pdf_dir=r"/data/fengyy/dateset/temp/0429/农村土地承包经营权登记簿/模板二(1)/模板二_002"
    # save_dir=r"/data/fengyy/dateset/temp/0429/result/模板二_002"
    # pdf_dir=r"/data/fengyy/dateset/temp/0429/农村土地承包经营权登记簿/模板二(1)/模板二_003"
    # save_dir=r"/data/fengyy/dateset/temp/0429/result/模板二_003"
    pdf_dir=r"/data/fengyy/dateset/temp/0429/农村土地承包经营权登记簿/模板一_1"
    save_dir=r"/data/fengyy/dateset/temp/0429/result/模板一_001"
    get_pdf(pdf_dir,save_dir)